/* Hybrid System Abstractor
 *
 * Copyright (c) 2013, Fortiss GmbH - An-Institut TU Muenchen
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.fortiss;

import java.util.*;
import com.microsoft.z3.*;
import org.fortiss.automaton.*;
import org.fortiss.util.Interval;
import org.fortiss.util.IntervalMatrix;
import org.fortiss.util.Matrix;
import org.fortiss.util.Predicate;

public class PredicateAbstractor
{
    private static int NUM_TERMS_MATRIX_EXPONENTIATION = 10;
    private static int NUM_INTERVALS = 2;
    
    private double T = 10;
    
    public KripkeStructure abstractWithPredicate(HybridSystem hs, List<List<Predicate>> excPred, Predicate initPred, int input) throws Exception
    {
        Context ctx = new Context();
        Set<Predicate> sp = new HashSet<Predicate>();
        for (List<Predicate> l: excPred)
            for (Predicate p: l)
                sp.add(p);
        List<Predicate> pred = new ArrayList<Predicate>(sp);
        Set<Set<Integer>> labels = getPossibleLabel(excPred, pred);
       
        KripkeStructure ks = new KripkeStructure();
        ks.init("predicate-abstracted " + hs.getName(), pred);
        
        Map<Integer, Map<Set<Integer>, KripkeNode>> node = new HashMap<Integer, Map<Set<Integer>, KripkeNode>>();
        int counter = 0;
        for (Node m: hs.getNode()){
            Map<Set<Integer>, KripkeNode> h = new HashMap<Set<Integer>, KripkeNode>();
            node.put(m.getID(), h);
            for (Set<Integer> l: labels){
                KripkeNode nd = new KripkeNode();
                
                nd.init(counter++, m.getName() + " [" + toString(pred, l, hs.getStateVariable()) + "]", checkInitial(ctx, m, pred, l, initPred), l);
                h.put(l, nd);
                ks.addNode(nd);
            }
        }
        
        for (Node src: hs.getNode()){
            HybridNode hsrc = (HybridNode) src;
            for (int i=0; i<input; i++){
                Set<ControlledUpdateEdge> ctrlE = hsrc.getControlledEdge(i);
                if (ctrlE.isEmpty()) ctrlE.add(null);
                for (ControlledUpdateEdge ce: ctrlE){
                    HybridNode hmid = ce==null ? hsrc : hs.getNode(ce.getDestinationID());
                    Set<AutonomousUpdateEdge> atnmE = hmid.getAutonomousEdge();
                    atnmE.add(null);
                    for (AutonomousUpdateEdge ae: atnmE){
                        HybridNode hdest = ae==null ? hmid : hs.getNode(ae.getDestinationID());
                        for (Set<Integer> lsrc: labels){
                            KripkeNode ksrc = node.get(hsrc.getID()).get(lsrc);
                            for (Set<Integer> ldest: labels){
                                KripkeNode kdest = node.get(hdest.getID()).get(ldest);
                                if (!ksrc.existEdge(i, kdest.getID()) && checkTransition(ctx, hs, hsrc, ks.toPredicate(lsrc), ce, ae, ks.toPredicate(ldest))){
                                    ControlledEdge ec = new ControlledEdge();
                                    ec.init(i, kdest.getID());
                                    ksrc.addEdge(ec);
                                }
                            }
                        }
                    }
                }
            }
        }
        
        return ks;
    }
    
    private Set<Set<Integer>> getPossibleLabel(List<List<Predicate>> llp, List<Predicate> p) throws Exception
    {
        return getPossibleLabel(llp, p, new ArrayList<Integer>(llp.size()));
    }
    
    private Set<Set<Integer>> getPossibleLabel(List<List<Predicate>> llp, List<Predicate> p, List<Integer> r) throws Exception
    {
        if (r.size()==llp.size()){
            Context ctx = new Context();
            BoolExpr e = ctx.mkBool(true);
            int n = llp.get(0).get(0).getVariable();
            e = ctx.mkAnd(e, generateDisjunctiveVariables(ctx, "x%d", n));
            IntervalMatrix mi = IntervalMatrix.eye(n);
            IntervalMatrix b  = IntervalMatrix.zero(n, 1);
            for (int i=0; i<r.size(); i++)
                e = ctx.mkAnd(e, generateWeakestConstraint(ctx, mi, b, llp.get(i).get(r.get(i)), "x%d"));
            
            Solver solver = ctx.mkSolver();
            solver.add(e);
            boolean check = solver.check()==Status.SATISFIABLE;
            solver.dispose();
            
            if (check){
                Set<Integer> s = new HashSet<Integer>();
                Set<Set<Integer>> ss = new HashSet<Set<Integer>>();
                for(int i=0; i<r.size(); i++) 
                    s.add(p.indexOf(llp.get(i).get(r.get(i))));
                ss.add(s);
                return ss;
            }
            return new HashSet<Set<Integer>>();
        }
        else {
            Set<Set<Integer>> ret = new HashSet<Set<Integer>>();
            for (int i=0; i<llp.get(r.size()).size(); i++){
                r.add(i);
                ret.addAll(getPossibleLabel(llp, p, r));
                r.remove(r.size()-1);
            }
            return ret;
        }
    }
    
    private boolean checkInitial(Context ctx, Node nd, List<Predicate> pred, Set<Integer> l, Predicate p) throws Exception
    {
        if (nd.isInitial()){
            BoolExpr e = ctx.mkBool(true);
            int n = p.getVariable();
            e = ctx.mkAnd(e, generateDisjunctiveVariables(ctx, "x%d", n));
            IntervalMatrix mi = IntervalMatrix.eye(n);
            IntervalMatrix b  = IntervalMatrix.zero(n, 1);
        
            for (Integer i: l)
                e = ctx.mkAnd(e, generateWeakestConstraint(ctx, mi, b, pred.get(i), "x%d"));
            
            e = ctx.mkAnd(e, generateWeakestConstraint(ctx, mi, b, p, "x%d"));
            
            Solver solver = ctx.mkSolver();
            solver.add(e);
            boolean ret = solver.check()==Status.SATISFIABLE;
            ctx.dispose();
            
            return ret;
        }
        else
            return false;
    }
    
    private String toString(List<Predicate> plist, Set<Integer> l, List<String> var)
    {
        String ret = "";
        for (Integer i: l){
            ret += ret.equals("") ? "" : " or ";
            ret += plist.get(i).toString(var);
        }
        return ret;
    }
    
    private boolean checkTransition(Context ctx, HybridSystem hs, HybridNode src, Set<Predicate> lsrc, ControlledUpdateEdge ce, AutonomousUpdateEdge ae, Set<Predicate> ldest) throws Exception
    {
        IntervalMatrix A = IntervalMatrix.eye(hs.getStateVariable().size());
        IntervalMatrix b = IntervalMatrix.zero(hs.getStateVariable().size(), 1);
        BoolExpr e = ctx.mkBool(true);

        e = ctx.mkAnd(e, generateDisjunctiveVariables(ctx, "x%d", hs.getStateVariable().size()));
        
        
        if (src.getInvariant()!=null)
            e = ctx.mkAnd(e, generateWeakestConstraint(ctx, A, b, src.getInvariant(), "x%d"));
        

        for (Predicate p: lsrc)
            e = ctx.mkAnd(e, generateWeakestConstraint(ctx, A, b, p, "x%d"));
        
        HybridNode from = src;
        if (ce!=null){
            from = hs.getNode(ce.getDestinationID());
            IntervalMatrix Au = ce.getUpdateFunction().delCoefficient();
            IntervalMatrix bu = ce.getUpdateFunction().getCoefficient();
            A = Au.times(A);
            b = Au.times(b).plus(bu);
        }
        
        if (ae==null){
            IntervalMatrix r12 = solveAffineDynamics(from.getFlow(), new Interval(T, T));
            
            A = r12.delCoefficient().times(A);
            b = r12.delCoefficient().times(b).plus(r12.getCoefficient());
                
            for (Predicate p: ldest)
                e = ctx.mkAnd(e, generateWeakestConstraint(ctx, A, b, p, "x%d"));
            
            Solver solver = ctx.mkSolver();
            solver.add(e);
            boolean ret = solver.check()==Status.SATISFIABLE;
            ctx.dispose();
            return ret;
        }
        else{
            double delta = T/NUM_INTERVALS;
            HybridNode dest = hs.getNode(ae.getDestinationID());
            
            for (int i=0; i<NUM_INTERVALS; i++){
                BoolExpr eSeg = e;
                IntervalMatrix r12 = solveAffineDynamics(from.getFlow(), new Interval(i*delta, (i+1)*delta));
                A = r12.delCoefficient().times(A);
                b = r12.delCoefficient().times(b).plus(r12.getCoefficient());
                
                eSeg = ctx.mkAnd(eSeg, generateWeakestConstraint(ctx, A, b, ae.getGuard(), "x%d"));

                IntervalMatrix r23 = solveAffineDynamics(dest.getFlow(), new Interval(T-(i+1)*delta, T-i*delta));
                A = r23.delCoefficient().times(A);
                b = r23.delCoefficient().times(b).plus(r23.getCoefficient());
                
                for (Predicate p: ldest)
                    eSeg = ctx.mkAnd(eSeg, generateWeakestConstraint(ctx, A, b, p, "x%d"));
                
                Solver solver = ctx.mkSolver();
                solver.add(eSeg);
                boolean ret = solver.check()==Status.SATISFIABLE;
                ctx.dispose();
                
                if (ret) return true;
            }
            return false;
        }
    }
    
    private IntervalMatrix solveAffineDynamics(IntervalMatrix affineDynamics, Interval time)
    {
        IntervalMatrix A = affineDynamics.delCoefficient();
        IntervalMatrix b = affineDynamics.getCoefficient();
        int n = A.getRow();
        IntervalMatrix TA = A.times(time);
        
        IntervalMatrix eTA = TA.expm();
        
        IntervalMatrix PAT = IntervalMatrix.eye(n).times(time);
        IntervalMatrix additionalTerm = PAT;
        for (int i = 1; i<= NUM_TERMS_MATRIX_EXPONENTIATION; i++) {
            additionalTerm = additionalTerm.times(A).times(time).times(1.0/(i+1));
            PAT = PAT.plus(additionalTerm);
        }
        
        return eTA.addCoefficient(PAT.times(b));
    }
    
    private BoolExpr generateDisjunctiveVariables(Context ctx, String t, int n) throws Exception
    {
        BoolExpr ret = ctx.mkBool(true);
        for (int i=0; i<n; i++){
            String var = String.format(t, i);
            String vp = var + "_p", vm = var + "_m";
            ret = ctx.mkAnd(ret, ctx.mkGe(ctx.mkRealConst(vp), ctx.mkReal(0)));
            ret = ctx.mkAnd(ret, ctx.mkGe(ctx.mkRealConst(vm), ctx.mkReal(0)));
            ret = ctx.mkAnd(ret, ctx.mkOr(ctx.mkEq(ctx.mkRealConst(vp), ctx.mkReal(0)), ctx.mkEq(ctx.mkRealConst(vm), ctx.mkReal(0))));
        }
        return ret;
    }
    
    private BoolExpr generateWeakestConstraint(Context ctx, IntervalMatrix A, IntervalMatrix b, Predicate p, String x) throws Exception
    {
        // predicate(A*x + b)
        BoolExpr ret = ctx.mkBool(true);
        if (p.getL()!=null){
            Matrix pL = p.getL();
            IntervalMatrix Apl = pL.delCoefficient().times(A);
            IntervalMatrix bpl = pL.delCoefficient().times(b);
            ret = ctx.mkAnd(ret, generateConstraint(ctx, Apl.getLower(), Apl.getUpper().times(-1), bpl.getLower(), pL.getCoefficient(), x+"_p", x+"_m", "<"));
        }
        if (p.getLE()!=null){
            Matrix pLE = p.getLE();
            IntervalMatrix Aple = pLE.delCoefficient().times(A);
            IntervalMatrix bple = pLE.delCoefficient().times(b);
            ret = ctx.mkAnd(ret, generateConstraint(ctx, Aple.getLower(), Aple.getUpper().times(-1), bple.getLower(), pLE.getCoefficient(), x+"_p", x+"_m", "<="));
        }
        return ret;
    }
    
    private BoolExpr generateConstraint(Context ctx, Matrix A1, Matrix A2, Matrix b1, Matrix b2, String x1, String x2, String op) throws Exception
    {
        // (A1*x1 + A2*x2 + b1) op b2
        BoolExpr ret = ctx.mkBool(true);
        for (int i=0; i<A1.getRow(); i++){
            ArithExpr e = ctx.mkReal(0);
            for (int j=0; j<A1.getColumn(); j++){
                e = ctx.mkAdd(e, ctx.mkMul(ctx.mkReal(String.format("%f", A1.getElement(i, j))), ctx.mkRealConst(String.format(x1, j))));
            }
            for (int j=0; j<A2.getColumn(); j++) {
                e = ctx.mkAdd(e, ctx.mkMul(ctx.mkReal(String.format("%f", A2.getElement(i, j))), ctx.mkRealConst(String.format(x2, j))));
            }
            e = ctx.mkAdd(e, ctx.mkReal(String.format("%f", b1.getElement(i, 0))));
            
            if (op.equals("<="))
                ret = ctx.mkAnd(ret, ctx.mkLe(e, ctx.mkReal(String.format("%f", b2.getElement(i, 0)))));
            else if (op.equals("<"))
                ret = ctx.mkAnd(ret, ctx.mkLt(e, ctx.mkReal(String.format("%f", b2.getElement(i, 0)))));
            else
                System.exit(1);
        }
        
        return ret;
    }
}
